#!/usr/bin/env python
'''
Given a fasta file containing many sequences, 
extract the longest CDS from the fasta.
'''

__author__ = "Yangho chen"
__email__ = "yanghoch@usc.edu"
__modification__ = "Created on Aug 16, 2009"

from parse_fasta_file import reverseCompliment
from parse_fasta_file import print_fasta
from parse_fasta_file import fasta_gen
#from parse_fasta_file import get_tag_seq_from_fasta_w_substr
import os
import getopt
import sys
import string

def isStopCoden ( mRNA, start ):
    codenLength = 3
    codon = string.join( mRNA[start:( start + codenLength )], "" )
    codon = string.upper( codon )
    if ( codon == "TAG" or
         codon == "TGA" or
         codon == "TAA" or
         codon == "UGA" or
         codon == "UAG" or
         codon == "UAA" ):
        return True
    else:
        return False

def isStartCoden ( mRNA, start ):
    codenLength = 3
    codon = string.join( mRNA[start:( start + codenLength )], "" )
    codon = string.upper( codon )
    if ( codon == "ATG" or
         codon == "TUG" ):
        return True
    else:
        return False

def check_frame( start, mRNA ):
    end = len( mRNA )
    ORF = ( end, end )
    i = start;
    codenLength = 3
    while ( i < end - codenLength + 1 ):
        if( isStartCoden( mRNA, i ) ):
            ORF = ( i, i )
            break
        i += codenLength

    while ( i < end - codenLength + 1 ):
        if( isStopCoden( mRNA, i ) ):
            ORF = ( ORF[0] , i + codenLength - 1 )
            break
        i += codenLength
    return ORF

def get_the_longest_ORF( mRNA ):
    '''return the longest ORF in a tuple with
    the two positions of the start and stop codon '''
    end = len( mRNA )
    longestORF = ( end, end )
    codenLength = 3
    for frame_shift in range ( codenLength ):
        ORF_search_Start = frame_shift
        ORF = ( end, end )
        while ( ORF_search_Start < end ):
            ORF = check_frame( ORF_search_Start, mRNA )
            if ORF[1] - ORF[0] > longestORF[1] - longestORF[0]:
                longestORF = ORF
            ORF_search_Start = ORF[1] + codenLength
    return longestORF

def print_CDSs_from_faFile( faFileN, oFileN ):
    ''' Given a .fasta file with many sequence,
    output another .fasta file with the longest ORFs seq
    corresponding to the seqs.'''
    if os.path.exists( oFileN ):
        os.remove( oFileN )
    # faFile = open( faFileN, "r" )
    seqFactory = fasta_gen( faFileN, '>' )
    seqNo = 0
    while 1:
        # [tag, seq] = get_tag_seq_from_fasta_w_substr( faFile, '>' )
        ( tag, seq ) = seqFactory.next()
        if tag != "":
            seqNo += 1
            ORF = get_the_longest_ORF( seq )
            # also check the reverse compliment string
            rc_seq = reverseCompliment( seq )
            rcORF = get_the_longest_ORF( rc_seq )

            ORF_Len = ORF[1] - ORF[0]
            rcORF_Len = rcORF[1] - rcORF[0]
            ORF_LENGTH_THRESHOLD = 300
            if( ( ORF_Len >= rcORF_Len ) and
                  ORF_Len > ORF_LENGTH_THRESHOLD ):
                tag = tag + ' ' + str( ORF_Len )
                print_fasta( oFileN, tag, seq[ORF[0]:( ORF[1] + 1 )] )
            elif ( ( rcORF_Len > ORF_Len ) and
                     rcORF_Len > ORF_LENGTH_THRESHOLD ):
                tag = "RC" + tag + str( rcORF_Len )
                print_fasta( oFileN, tag, rc_seq[rcORF[0]:( rcORF[1] + 1 )] )
        else :
            # faFile.close()
            return seqNo

def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print "usage: ", sys.argv[0]
    print "The two parameters should be "
    print "(1) fasta file"
    print "(2) The output file name"
    sys.exit( 2 )

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error:
        usage()
        sys.exit( 2 );
    # process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            usage()
            sys.exit();
    if len( args ) == 2 :
        faFileN = args[0]
        ofileN = args[1]
        seqNo = print_CDSs_from_faFile( faFileN, ofileN )
        print "There are", seqNo, "sequences."
    else:
        usage()

if __name__ == '__main__':
    main()
